package developer.zyk.ability.config;

import developer.zyk.ability.chat.ability.memory.msgHandler.MessageHandler;
import developer.zyk.ability.chat.ability.rag.vector.MysqlVectorStore;
import developer.zyk.ability.knowledge.dao.mappers.ChatKnowledgeRecordDetailMapper;
import developer.zyk.ability.knowledge.dao.mappers.ChatKnowledgeRecordMapper;
import developer.zyk.global.common.consts.ChatRole;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.client.advisor.SimpleLoggerAdvisor;
import org.springframework.ai.embedding.EmbeddingModel;
import org.springframework.ai.openai.OpenAiChatModel;
import org.springframework.ai.vectorstore.SimpleVectorStore;
import org.springframework.ai.vectorstore.VectorStore;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.beans.factory.annotation.Qualifier;
import org.springframework.boot.context.properties.ConfigurationPropertiesScan;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;

import java.util.HashMap;
import java.util.Map;

@Configuration
@ConfigurationPropertiesScan({"developer.zyk.ability.config"})
public class AIConfig {

    @Autowired
    @Qualifier("MessageHandlerUser")
    private MessageHandler messageHandlerUser;

    @Autowired
    @Qualifier("messageHandlerAssistant")
    private MessageHandler messageHandlerAssistant;

    @Bean
    public ChatClient chatClient(OpenAiChatModel openAiChatModel) {
        return ChatClient.builder(openAiChatModel)
                .defaultAdvisors(new SimpleLoggerAdvisor())
                .build();
    }

    @Bean("messageConverter")
    public Map<ChatRole, MessageHandler> messageConverter() {
        Map<ChatRole, MessageHandler> map = new HashMap<>();
        map.put(ChatRole.USER, messageHandlerUser);
        map.put(ChatRole.ASSISTANT, messageHandlerAssistant);
        return map;
    }

    @Bean("vectorStore")
    public VectorStore vectorStore(
            @Autowired @Qualifier("aliEmbeddingModel") EmbeddingModel embeddingModel,
            @Autowired ChatKnowledgeRecordDetailMapper chatKnowledgeRecordDetailMapper,
            @Autowired ChatKnowledgeRecordMapper chatKnowledgeRecordMapper
    ) {

        return new MysqlVectorStore(SimpleVectorStore.builder(embeddingModel),
                chatKnowledgeRecordDetailMapper,
                chatKnowledgeRecordMapper
        );
    }

}
